{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Minimize Malware Spread II"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #depth-first-search #breadth-first-search #union-find #graph #hash-table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #深度优先搜索 #广度优先搜索 #并查集 #图 #哈希表"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: minMalwareSpread"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #尽量减少恶意软件的传播 II"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给定一个由 <code>n</code> 个节点组成的网络，用 <code>n x n</code> 个邻接矩阵&nbsp;<code>graph</code>&nbsp;表示。在节点网络中，只有当&nbsp;<code>graph[i][j] = 1</code>&nbsp;时，节点&nbsp;<code>i</code>&nbsp;能够直接连接到另一个节点&nbsp;<code>j</code>。</p>\n",
    "\n",
    "<p>一些节点&nbsp;<code>initial</code>&nbsp;最初被恶意软件感染。只要两个节点直接连接，且其中至少一个节点受到恶意软件的感染，那么两个节点都将被恶意软件感染。这种恶意软件的传播将继续，直到没有更多的节点可以被这种方式感染。</p>\n",
    "\n",
    "<p>假设&nbsp;<code>M(initial)</code>&nbsp;是在恶意软件停止传播之后，整个网络中感染恶意软件的最终节点数。</p>\n",
    "\n",
    "<p>我们可以从 <code>initial</code> 中<strong>删除一个节点</strong>，<strong>并完全移除该节点以及从该节点到任何其他节点的任何连接。</strong></p>\n",
    "\n",
    "<p>请返回移除后能够使&nbsp;<code>M(initial)</code>&nbsp;最小化的节点。如果有多个节点满足条件，返回索引 <strong>最小的节点</strong> 。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<ol>\n",
    "</ol>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>graph = [[1,1,0],[1,1,0],[0,0,1]], initial = [0,1]\n",
    "<strong>输出：</strong>0\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>graph = [[1,1,0],[1,1,1],[0,1,1]], initial = [0,1]\n",
    "<strong>输出：</strong>1\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 3：</strong></p>\n",
    "\n",
    "<pre>\n",
    "<strong>输入：</strong>graph = [[1,1,0,0],[1,1,1,0],[0,1,1,1],[0,0,1,1]], initial = [0,1]\n",
    "<strong>输出：</strong>1\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "<meta charset=\"UTF-8\" />\n",
    "\n",
    "<ul>\n",
    "\t<li><code>n == graph.length</code></li>\n",
    "\t<li><code>n == graph[i].length</code></li>\n",
    "\t<li><code>2 &lt;= n &lt;= 300</code></li>\n",
    "\t<li><code>graph[i][j]</code>&nbsp;是&nbsp;<code>0</code>&nbsp;或&nbsp;<code>1</code>.</li>\n",
    "\t<li><code>graph[i][j] == graph[j][i]</code></li>\n",
    "\t<li><code>graph[i][i] == 1</code></li>\n",
    "\t<li><code>1 &lt;= initial.length &lt;&nbsp;n</code></li>\n",
    "\t<li><code>0 &lt;= initial[i] &lt;= n - 1</code></li>\n",
    "\t<li>&nbsp;<code>initial</code>&nbsp;中每个整数都<strong>不同</strong></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [minimize-malware-spread-ii](https://leetcode.cn/problems/minimize-malware-spread-ii/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [minimize-malware-spread-ii](https://leetcode.cn/problems/minimize-malware-spread-ii/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[[1,1,0],[1,1,0],[0,0,1]]\\n[0,1]', '[[1,1,0],[1,1,1],[0,1,1]]\\n[0,1]', '[[1,1,0,0],[1,1,1,0],[0,1,1,1],[0,0,1,1]]\\n[0,1]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        # 1. build union find disjoint set without vertices in the initial array\n",
    "        n = len(graph)\n",
    "        initial = set(initial)\n",
    "        UF = UnionFind(n)\n",
    "\n",
    "        for i in range(n):\n",
    "            if i not in initial:\n",
    "                for j in range(n):\n",
    "                    if (j not in initial) and graph[i][j]:\n",
    "                        UF.unionSet(i, j)\n",
    "        \n",
    "        # 2. find root of each vertex in the initial array\n",
    "        hashmap = dict()\n",
    "        for i in initial:\n",
    "            roots = set([UF.findSet(j) for j in range(n) if (j not in initial) and graph[i][j]])\n",
    "            hashmap[i] = roots \n",
    "        \n",
    "        rootCounter = collections.Counter(itertools.chain(*hashmap.values()))\n",
    "\n",
    "\n",
    "        # 3. compare max reduced number of infected of each vertex in the initial array\n",
    "        initial = sorted(list(initial))\n",
    "        ans = initial[0]\n",
    "        maxSize = -1\n",
    "        for i in initial:\n",
    "            size = sum([UF.sizeSets[root] for root in hashmap[i] if rootCounter[root] == 1])\n",
    "            if size > maxSize:\n",
    "                ans = i \n",
    "                maxSize = size \n",
    "        return ans \n",
    "\n",
    "\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int, m: int=None) -> None:\n",
    "        \"\"\"\n",
    "        build m disjoint sets from n nodes\n",
    "        T: O(N)\n",
    "        \"\"\"\n",
    "        self.parents = list(range(n))  \n",
    "        self.ranks = [0] * n            \n",
    "        self.sizeSets = [1] * n         # sizeSets[i]: size of disjoint set rooted at vertex i\n",
    "\n",
    "    \n",
    "    def findSet(self, i: int) -> int:\n",
    "        \"\"\"return root of node i's set\n",
    "           T: amortized O(α(N))  effectively O(1)\n",
    "             first O(N) then O(1)\n",
    "           path-compression  \n",
    "           Recursion\n",
    "        \"\"\"\n",
    "        if self.parents[i] == i:\n",
    "            return i\n",
    "        else:\n",
    "            self.parents[i] = self.findSet(self.parents[i])\n",
    "            return self.parents[i]\n",
    "        \n",
    "\n",
    "    \n",
    "    def unionSet(self, i: int, j: int) -> None:\n",
    "        \"\"\"若节点i和节点j来自2个不同的互斥集合，将两个集合合并 \n",
    "           union by rank   矮树作为高树的子树\n",
    "           T: amortized O(α(N))  effectively O(1)\n",
    "             first O(N) then O(1)\n",
    "        \"\"\"\n",
    "        ip, jp = self.findSet(i), self.findSet(j)\n",
    "\n",
    "        # 若两个节点来自同一互斥集合，返回\n",
    "        if ip == jp:\n",
    "            return \n",
    "\n",
    "        # 若A树矮，B树高，A树连到B树的根节点上，新树rank不变\n",
    "        if self.ranks[ip] < self.ranks[jp]:\n",
    "            self.parents[ip] = jp\n",
    "            self.sizeSets[jp] += self.sizeSets[ip]\n",
    "\n",
    "        # 若B树矮，A树高，B树连到A树的根节点上，新树rank不变\n",
    "        elif self.ranks[jp] < self.ranks[ip]:\n",
    "            self.parents[jp] = ip\n",
    "            self.sizeSets[ip] += self.sizeSets[jp]\n",
    "\n",
    "        # 若两树等高，任选一种方案，新树 rank +1\n",
    "        else:\n",
    "            self.parents[jp] = ip\n",
    "            self.sizeSets[ip] += self.sizeSets[jp]\n",
    "            self.ranks[ip] += 1\n",
    "\n",
    "    \n",
    "    def isSameSet(self, i: int, j: int) -> bool:\n",
    "        \"\"\"判断节点i和节点j是否属于同一集合\n",
    "           T: O(α(N) first O(N) then O(1)\n",
    "        \"\"\"\n",
    "        # path compression twice\n",
    "        return self.findSet(i) == self.findSet(j)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parent = list(range(n))\n",
    "    \n",
    "    def union(self, index1, index2):\n",
    "        self.parent[self.find(index2)] = self.find(index1)\n",
    "\n",
    "    def find(self, index):\n",
    "        if self.parent[index] != index:\n",
    "            self.parent[index] = self.find(self.parent[index])\n",
    "        return self.parent[index]\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n, ni = len(graph), len(initial)\n",
    "        res = 0\n",
    "        minnum = 1000\n",
    "        for e in initial:\n",
    "            union = UnionFind(n)\n",
    "            for i in range(n):\n",
    "                for j in range(i):\n",
    "                    if i != e and j != e and graph[i][j]:\n",
    "                        union.union(i, j)\n",
    "            for i in range(ni):\n",
    "                for j in range(i):\n",
    "                    if initial[i] != e and initial[j] != e:\n",
    "                        union.union(initial[i], initial[j])\n",
    "            group = initial[0] if initial[0] != e else initial[-1]\n",
    "            temp = 0\n",
    "            for i in range(n):\n",
    "                if union.find(i) == union.find(group):\n",
    "                    temp += 1\n",
    "            if temp < minnum or (temp == minnum and e < res):\n",
    "                minnum = temp\n",
    "                res = e\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "\n",
    "        def find(parent, x):\n",
    "            if parent[x] == x: return x\n",
    "            parent[x] = find(parent, parent[x])\n",
    "            return parent[x]\n",
    "\n",
    "        def func(skip):\n",
    "            parent = [i for i in range(n)]\n",
    "            size = [1] * n\n",
    "\n",
    "            for i in range(n):\n",
    "                for j in range(i + 1, n):\n",
    "                    if i == skip or j == skip: continue\n",
    "                    if graph[i][j] == 0: continue\n",
    "                    r0, r1 = find(parent, i), find(parent, j)\n",
    "                    if r0 != r1:\n",
    "                        parent[r1] = r0\n",
    "                        size[r0] += size[r1]\n",
    "            ans = 0\n",
    "            for v in initial:\n",
    "                if v == skip: continue\n",
    "                r = find(parent, v)\n",
    "                ans += size[r]\n",
    "                size[r] = 0\n",
    "            return ans\n",
    "\n",
    "        return min([[func(i), i] for i in initial])[1]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self.root = [i for i in range(n)]\n",
    "        self.size = [1] * n\n",
    "        self.part = n\n",
    "        return\n",
    "\n",
    "    def find(self, x):\n",
    "        lst = []\n",
    "        while x != self.root[x]:\n",
    "            lst.append(x)\n",
    "            # 在查询的时候合并到顺带直接根节点\n",
    "            x = self.root[x]\n",
    "        for w in lst:\n",
    "            self.root[w] = x\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        root_x = self.find(x)\n",
    "        root_y = self.find(y)\n",
    "        if root_x == root_y:\n",
    "            return False\n",
    "        if self.size[root_x] >= self.size[root_y]:\n",
    "            root_x, root_y = root_y, root_x\n",
    "        self.root[root_x] = root_y\n",
    "        self.size[root_y] += self.size[root_x]\n",
    "        # 将非根节点的秩赋0\n",
    "        self.size[root_x] = 0\n",
    "        self.part -= 1\n",
    "        return True\n",
    "\n",
    "    def is_connected(self, x, y):\n",
    "        return self.find(x) == self.find(y)\n",
    "\n",
    "    def get_root_part(self):\n",
    "        # 获取每个根节点对应的组\n",
    "        part = defaultdict(list)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            part[self.find(i)].append(i)\n",
    "        return part\n",
    "\n",
    "    def get_root_size(self):\n",
    "        # 获取每个根节点对应的组大小\n",
    "        size = defaultdict(int)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            size[self.find(i)] = self.size[self.find(i)]\n",
    "        return size\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        ans = -1\n",
    "        res = inf\n",
    "        for i in initial:\n",
    "            uf = UnionFind(n)\n",
    "            dct = set(initial)\n",
    "            dct.discard(i)\n",
    "            for x in range(n):\n",
    "                for y in range(x+1, n):\n",
    "                    if x != i and y != i and graph[x][y]:\n",
    "                        uf.union(x, y)\n",
    "            group = uf.get_root_part()\n",
    "            cur = 0\n",
    "            for g in group:\n",
    "                if any(x in dct for x in group[g]):\n",
    "                    cur += len(group[g])\n",
    "            if cur < res or (cur==res and i<ans):\n",
    "                ans = i\n",
    "                res = cur\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        def bfs(remove: int) -> int:\n",
    "            s = 0\n",
    "            q = deque(x for x in initial if x != remove)\n",
    "            vis = [False] * n\n",
    "            for x in initial:\n",
    "                if x != remove:\n",
    "                    vis[x] = True\n",
    "            while q:\n",
    "                p = q.popleft()\n",
    "                s += 1\n",
    "                for i in range(n):\n",
    "                    if graph[p][i] and i != remove and not vis[i]:\n",
    "                        vis[i] = True\n",
    "                        q.append(i)\n",
    "            return s\n",
    "        res = {}\n",
    "        for remove in initial:\n",
    "            res[remove] = bfs(remove)\n",
    "        mn = min(res.values())\n",
    "        res = [k for k in res if res[k] == mn]\n",
    "        return min(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        m, n = len(graph), len(graph[0])\n",
    "        def check_num(x: int) -> int:\n",
    "            vis = set([y for y in initial if y != x])\n",
    "            ini = deque([y for y in vis])\n",
    "            while ini:\n",
    "                tmp = ini.popleft()\n",
    "                for i in range(n):\n",
    "                    if i != x and graph[tmp][i] == 1 and i not in vis:\n",
    "                        vis.add(i)\n",
    "                        ini.append(i)\n",
    "            return len(vis)\n",
    "        ans = [inf, inf]\n",
    "        for x in initial:\n",
    "            if check_num(x) < ans[0]:\n",
    "                ans = [check_num(x), x]\n",
    "            elif check_num(x) == ans[0] and x < ans[1]:\n",
    "                ans = [check_num(x), x]\n",
    "        return ans[1]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "            \n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    @staticmethod\n",
    "    def myFind(parent, i):\n",
    "        if parent[i] != i: \n",
    "            parent[i] = Solution.myFind(parent, parent[i])\n",
    "        return parent[i]\n",
    "    \n",
    "    @staticmethod\n",
    "    def myUnion(parent, rank, i, j):\n",
    "        pi, pj = Solution.myFind(parent, i), Solution.myFind(parent, j)\n",
    "        if pi == pj: return\n",
    "        if rank[pi] > rank[pj]:\n",
    "            parent[pj] = pi\n",
    "        elif rank[pi] < rank[pj]:\n",
    "            parent[pi] = pj\n",
    "        else:\n",
    "            parent[pj] = pi\n",
    "            rank[pi] += 1\n",
    "        return\n",
    "\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial.sort()\n",
    "        def infect(delete):\n",
    "            n = len(graph)\n",
    "            parent, rank = [_ for _ in range(n)], [0] * n\n",
    "            for i in range(n):\n",
    "                for j in range(i):\n",
    "                    if i!= delete and j!= delete and graph[i][j] == 1:\n",
    "                        Solution.myUnion(parent, rank, i, j)    \n",
    "            for i in range(n): _ = Solution.myFind(parent, i)\n",
    "            count = Counter(parent)\n",
    "            s = set()\n",
    "            for i in initial:\n",
    "                if i != delete: \n",
    "                    s.add(parent[i])\n",
    "            ans = 0\n",
    "            for i in s:\n",
    "                ans += count[i]\n",
    "            return ans\n",
    "        tmp = len(graph)\n",
    "        ans = initial[0]\n",
    "        for i in initial:\n",
    "            tmp1 = infect(i)\n",
    "            if tmp1 < tmp:\n",
    "                ans = i\n",
    "                tmp = tmp1\n",
    "        return ans\n",
    "\n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        def bfs(remove: int) -> int:\n",
    "            s = 0\n",
    "            q = deque(x for x in initial if x != remove)\n",
    "            vis = [False] * n\n",
    "            for x in initial:\n",
    "                if x != remove:\n",
    "                    vis[x] = True\n",
    "            while q:\n",
    "                p = q.popleft()\n",
    "                s += 1\n",
    "                for i in range(n):\n",
    "                    if graph[p][i] and i != remove and not vis[i]:\n",
    "                        vis[i] = True\n",
    "                        q.append(i)\n",
    "            return s\n",
    "        res = {}\n",
    "        for remove in initial:\n",
    "            res[remove] = bfs(remove)\n",
    "        mn = min(res.values())\n",
    "        res = [k for k in res if res[k] == mn]\n",
    "        return min(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "pa = [0] * 300\r\n",
    "sz = [0] * 300\r\n",
    "\r\n",
    "def find(i: int) -> int:\r\n",
    "    if i != pa[i]:\r\n",
    "        pa[i] = find(pa[i])\r\n",
    "    return pa[i]\r\n",
    "\r\n",
    "def join(i: int, j: int) -> None:\r\n",
    "    f, g = find(i), find(j)\r\n",
    "    if f == g: return\r\n",
    "    if sz[f] > sz[g]:\r\n",
    "        sz[f] += sz[g]\r\n",
    "        pa[g] = f\r\n",
    "    else:\r\n",
    "        sz[g] += sz[f]\r\n",
    "        pa[f] = g\r\n",
    "    return\r\n",
    "\r\n",
    "class Solution:\r\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\r\n",
    "        n = len(graph)\r\n",
    "        for i in range(n):\r\n",
    "            pa[i] = i\r\n",
    "            sz[i] = 1\r\n",
    "        \r\n",
    "        nonInit = set(range(n)).difference(initial)\r\n",
    "        for i in nonInit:\r\n",
    "            for j in nonInit:\r\n",
    "                if graph[i][j] == 1:\r\n",
    "                    join(i, j)\r\n",
    "\r\n",
    "        infectedBy: dict[int, set[int]] = {i: set() for i in map(find, nonInit)}\r\n",
    "        for i in nonInit:\r\n",
    "            for j in initial:\r\n",
    "                if graph[i][j] == 1:\r\n",
    "                    infectedBy[find(i)].add(j)\r\n",
    "        \r\n",
    "        effectNodes: dict[int, int] = dict()\r\n",
    "        for root, inits in infectedBy.items():\r\n",
    "            if len(inits) == 1:\r\n",
    "                init = inits.pop()\r\n",
    "                effectNodes.setdefault(init, 0)\r\n",
    "                effectNodes[init] += sz[root]\r\n",
    "        \r\n",
    "        if not effectNodes: return min(initial)\r\n",
    "        return min(effectNodes.keys(), key=lambda node: (-effectNodes[node], node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial.sort()\n",
    "        ans, ans_v = inf, inf\n",
    "        for ind, init in enumerate(initial) :\n",
    "            tmp = initial[:ind] + initial[ind+1:]\n",
    "            visit = set(tmp)\n",
    "            q = list(tmp)\n",
    "            while q :\n",
    "                node = q.pop()\n",
    "                for i, v in enumerate(graph[node]) :\n",
    "                    if v and i != init and i not in visit:\n",
    "                        visit.add(i)\n",
    "                        q.append(i)\n",
    "            if len(visit) < ans_v:\n",
    "                ans = init\n",
    "                ans_v = len(visit)\n",
    "        \n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial = set(initial)\n",
    "        def find(x):#并查集\n",
    "            if x not in cache:\n",
    "                cache[x] = x\n",
    "            if cache[x] != x:\n",
    "                cache[x] = find(cache[x])\n",
    "            return cache[x]\n",
    "        def merge(x, y):\n",
    "            cache[find(x)] = find(y)\n",
    "        n = len(graph)\n",
    "        cache = {}\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if i != j and i not in initial and j not in initial and graph[i][j] == 1:\n",
    "                    merge(i, j)\n",
    "        d = {}#计算每个节点集合的数量\n",
    "        for i in range(n):\n",
    "            i = find(i)\n",
    "            d[i] = d.get(i, 0) + 1\n",
    "        s = set(d.keys())\n",
    "        cnt = {}#存储每个节点集合能被多少个感染节点所连接\n",
    "        for i in initial:\n",
    "            c = 0\n",
    "            s1 = set([find(j) for j in range(n) if j not in initial and graph[i][j] == 1])\n",
    "            for j in s1:\n",
    "                cnt[j] = cnt.get(j, 0) + 1\n",
    "        m = -1\n",
    "        for i in sorted(initial):#升序遍历,如果有多个节点需要返回最小的\n",
    "            c = 0\n",
    "            s1 = set([find(j) for j in range(n) if j not in initial and graph[i][j] == 1])\n",
    "            for j in s1:\n",
    "                if cnt[j] == 1:#如果节点集合值为1,说明去除本节点后这个集合不会被感染,进行累加计数\n",
    "                    c += d[j]\n",
    "            if c > m:#当不会被感染的集合数量越大说明被感染节点总数越小\n",
    "                m = c\n",
    "                res = i\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial = set(initial)\n",
    "        n = len(graph)\n",
    "        clean = [i for i in range(n) if i not in initial]\n",
    "\n",
    "        p = [i for i in range(n)]\n",
    "        s = [1] * n \n",
    "\n",
    "        def find(x):\n",
    "            if x != p[x]:\n",
    "                p[x] = find(p[x])\n",
    "                # x = p[x]\n",
    "            return p[x]\n",
    "        \n",
    "        def union(x, y):\n",
    "            rx, ry = find(x), find(y)\n",
    "            if rx != ry:\n",
    "                p[ry] = rx\n",
    "                s[rx] += s[ry]\n",
    "        \n",
    "        for i in clean:\n",
    "            for j in clean:\n",
    "                if graph[i][j]:\n",
    "                    union(i, j)\n",
    "\n",
    "        cnt = Counter()\n",
    "        node = {}\n",
    "        for i in initial:\n",
    "            st = set()\n",
    "            for j in clean:\n",
    "                if graph[i][j]:\n",
    "                    st.add(find(j))\n",
    "            node[i] = st\n",
    "            for c in st:\n",
    "                cnt[c] += 1\n",
    "        \n",
    "        best = (-1, None)\n",
    "        for u, vs in node.items():\n",
    "            score = 0\n",
    "            for c in vs:\n",
    "                if cnt[c] == 1:\n",
    "                    score += s[c]\n",
    "            if score > best[0] or score == best[0] and u < best[1]:\n",
    "                best = (score, u)\n",
    "        \n",
    "        return best[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class DSU:\n",
    "    def __init__(self, n):\n",
    "        self.states = [i for i in range(n)]\n",
    "\n",
    "    def find(self, i):\n",
    "        if self.states[i] != i: self.states[i] = self.find(self.states[i])\n",
    "        return self.states[i]\n",
    "\n",
    "    def union(self, i, j):\n",
    "        self.states[self.find(i)] = self.find(j)\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n, ans, mx = len(graph), -1, -1\n",
    "        for idx in initial:\n",
    "            dsu = DSU(n)\n",
    "            for i in range(n):\n",
    "                if i == idx: continue\n",
    "                for j in range(i + 1, n):\n",
    "                    if j == idx: continue\n",
    "                    if graph[i][j] == 1: dsu.union(i, j)\n",
    "            cnts = collections.Counter([dsu.find(i) for i in range(n)])\n",
    "            for i in initial:\n",
    "                if i != idx: cnts[dsu.find(i)] = 0\n",
    "            tmp = sum([cnts[k] for k in cnts])\n",
    "            if tmp > mx:\n",
    "                mx, ans = tmp, idx\n",
    "            elif tmp == mx and idx < ans: ans = idx\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    MAXN = 301\n",
    "\n",
    "    def __init__(self):\n",
    "        self.virus = [False] * self.MAXN\n",
    "        self.cnts = [0] * self.MAXN\n",
    "        self.infect = [-1] * self.MAXN\n",
    "        self.father = [i for i in range(self.MAXN)]\n",
    "        self.size = [1] * self.MAXN\n",
    "\n",
    "    def build(self, n, initial):\n",
    "        # 初始化并查集和病毒信息\n",
    "        for i in range(n):\n",
    "            self.virus[i] = False\n",
    "            self.cnts[i] = 0\n",
    "            self.infect[i] = -1\n",
    "            self.size[i] = 1\n",
    "            self.father[i] = i\n",
    "        for i in initial:\n",
    "            self.virus[i] = True\n",
    "\n",
    "    def find(self, i):\n",
    "        # 寻找元素 i 所在集合的代表节点\n",
    "        if i != self.father[i]:\n",
    "            self.father[i] = self.find(self.father[i])\n",
    "        return self.father[i]\n",
    "\n",
    "    def union(self, x, y):\n",
    "        # 合并两个集合\n",
    "        fx = self.find(x)\n",
    "        fy = self.find(y)\n",
    "        if fx != fy:\n",
    "            self.father[fx] = fy\n",
    "            self.size[fy] += self.size[fx]\n",
    "\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        self.build(n, initial)\n",
    "        \n",
    "        # 不是病毒的点，普通点合并！\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j] == 1 and not self.virus[i] and not self.virus[j]:\n",
    "                    self.union(i, j)\n",
    "        \n",
    "        # 病毒周围的普通点(集合 )去设置源头！\n",
    "        for sick in initial:\n",
    "            for neighbor in range(n):\n",
    "                if sick != neighbor and not self.virus[neighbor] and graph[sick][neighbor] == 1:\n",
    "                    fn = self.find(neighbor)\n",
    "                    if self.infect[fn] == -1:\n",
    "                        self.infect[fn] = sick\n",
    "                    elif self.infect[fn] != -2 and self.infect[fn] != sick:\n",
    "                        self.infect[fn] = -2\n",
    "        \n",
    "        # 统计拯救数据\n",
    "        for i in range(n):\n",
    "            # 不是代表点，不看\n",
    "            if i == self.find(i) and self.infect[i] >= 0:\n",
    "                self.cnts[self.infect[i]] += self.size[i]\n",
    "        \n",
    "        # 找到最终结果\n",
    "        initial.sort()\n",
    "        ans = initial[0]\n",
    "        max_cnt = self.cnts[ans]\n",
    "        for i in initial:\n",
    "            if self.cnts[i] > max_cnt:\n",
    "                ans = i\n",
    "                max_cnt = self.cnts[i]\n",
    "        \n",
    "        return ans\n",
    "\n",
    "# 测试用例\n",
    "solution = Solution()\n",
    "\n",
    "# 替换为你的测试数据\n",
    "graph = [[1, 1, 0], [1, 1, 0], [0, 0, 1]]\n",
    "initial = [0, 1]\n",
    "\n",
    "print(solution.minMalwareSpread(graph, initial))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from math import *\n",
    "from queue import PriorityQueue,Queue\n",
    "from typing import *\n",
    "from heapq import *\n",
    "from functools import lru_cache,cache\n",
    "from string import *\n",
    "# by hangpengjie\n",
    "class UF:\n",
    "    def __init__(self,n:int) -> None:\n",
    "        self.parent = [i for i in range(n)]\n",
    "        self.rank = [1 for _ in range(n)]\n",
    "    def find(self,x:int)->int:\n",
    "        if self.parent[x] == x:\n",
    "            return x\n",
    "        self.parent[x] = self.find(self.parent[x])\n",
    "        return self.parent[x]\n",
    "    def merge(self,x:int,y:int)->None:\n",
    "        x = self.find(x)\n",
    "        y = self.find(y)\n",
    "        if x == y:\n",
    "            return\n",
    "        self.rank[y] += self.rank[x]\n",
    "        self.rank[x] = 0\n",
    "        self.parent[x] = y\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        clean = [1] * n\n",
    "        # 标记出受污染的软件\n",
    "        for t in initial:\n",
    "            clean[t] = 0\n",
    "        uf = UF(n)\n",
    "        # 建立未被污染的软件的并查集\n",
    "        for i in range(n):\n",
    "            if clean[i]:\n",
    "                for j in range(n):\n",
    "                    if clean[j]:\n",
    "                        if graph[i][j]:\n",
    "                            uf.merge(i,j)\n",
    "        # 统计并查集中每个集合和哪些受污染的软件连接\n",
    "        mp = dict()\n",
    "        for i in range(n):\n",
    "            if not clean[i]:\n",
    "                for j in range(n):\n",
    "                    if clean[j] and graph[i][j]:\n",
    "                        p = uf.find(j)\n",
    "                        if p not in mp:\n",
    "                            mp[p] = set()\n",
    "                        mp[p].add(i)\n",
    "        # 只记录集合只和一个受污染软件连接的受污染软件能够贡献的答案，即集合的秩\n",
    "        d = dict()\n",
    "        for k,v in mp.items():\n",
    "            if len(v) == 1:\n",
    "                t = v.pop()\n",
    "                if t not in d:\n",
    "                    d[t] = 0\n",
    "                d[t] += uf.rank[uf.find(k)]\n",
    "        initial.sort()\n",
    "        # 初始化答案\n",
    "        res = initial[0]\n",
    "        cnt = 0\n",
    "        for k,v in d.items():\n",
    "            if v == cnt and k < res:\n",
    "                res = k\n",
    "            elif v > cnt:\n",
    "                cnt = v\n",
    "                res = k\n",
    "        return res\n",
    "    \n",
    "s = Solution()\n",
    "s.minMalwareSpread([[1,0,0,0,0,0,0,0,0],[0,1,0,0,0,0,0,0,1],[0,0,1,0,0,0,0,0,0],[0,0,0,1,0,0,0,0,1],[0,0,0,0,1,0,1,1,1],[0,0,0,0,0,1,0,0,1],[0,0,0,0,1,0,1,1,0],[0,0,0,0,1,0,1,1,0],[0,1,0,1,1,1,0,0,1]],\n",
    "                   [8,4,2,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        size = [1] * n\n",
    "        cnts = [0] * n\n",
    "        fa = list(range(n))\n",
    "        infect = [-1] * n\n",
    "        virus = [False] * n\n",
    "        for x in initial:\n",
    "            virus[x] = True\n",
    "        def find(x: int)-> int:\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(x:int, y:int)-> None:\n",
    "            fa_x = find(x)\n",
    "            fa_y = find(y)\n",
    "            if fa_x != fa_y:\n",
    "                fa[fa_x] = fa_y\n",
    "                size[fa_y] += size[fa_x]\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j] == 1 and not virus[i] and not virus[j]:\n",
    "                    union(i, j)\n",
    "        for sick in initial:\n",
    "            for i in range(n):\n",
    "                if sick != i and not virus[i] and graph[sick][i] == 1:\n",
    "                    fi = find(i)\n",
    "                    if infect[fi] == -1:\n",
    "                        infect[fi] = sick\n",
    "                    elif infect[fi] != -2 and infect[fi] != sick:\n",
    "                        infect[fi] = -2\n",
    "        for i in range(n):\n",
    "            if i == find(i) and infect[i] >= 0:\n",
    "                cnts[infect[i]] += size[i]\n",
    "        initial.sort()\n",
    "        ans = initial[0]\n",
    "        maxcnt = cnts[ans]\n",
    "        for i in initial:\n",
    "            if cnts[i] > maxcnt:\n",
    "                ans = i\n",
    "                maxcnt = cnts[i]\n",
    "        return ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        edges = defaultdict(list)\n",
    "        n = len(graph)\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j]:\n",
    "                    edges[i].append(j)\n",
    "                    edges[j].append(i)\n",
    "        cnt, ansNode, ansIdx = float(\"inf\"), None, None\n",
    "        for i in range(len(initial)):\n",
    "            q = initial[:i] + initial[i+1:]\n",
    "            q = deque(q)\n",
    "            curNode = initial[i]\n",
    "            visited = {node:True for node in q}\n",
    "            while q:\n",
    "                node = q.popleft()\n",
    "                for nextNode in edges[node]:\n",
    "                    if nextNode == curNode or nextNode in visited:\n",
    "                        continue\n",
    "                    q.append(nextNode)\n",
    "                    visited[nextNode] = True\n",
    "            curCnt = len(visited.keys())\n",
    "            # print(f\"curCnt:{curCnt}, curAns:{curNode}\")\n",
    "            if cnt > curCnt or (cnt == curCnt and ansNode > curNode):\n",
    "                cnt = curCnt\n",
    "                ansNode = curNode\n",
    "            # print(f\"cnt:{cnt}, ansNode:{ansNode}\")\n",
    "        return ansNode\n",
    "\n",
    "                    \n",
    "                    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n=len(graph)\n",
    "        def dfs(x):\n",
    "            find=[False]*n\n",
    "            q=[]\n",
    "            for i in initial:\n",
    "                find[i]=True\n",
    "                if i==x:\n",
    "                    continue\n",
    "                else:\n",
    "                    q.append(i)\n",
    "                    \n",
    "            c=0\n",
    "            while q:\n",
    "                tp=q.pop()\n",
    "                #if find[tp]:\n",
    "                #    continue\n",
    "                c+=1\n",
    "                #find[tp]=True\n",
    "                for i in range(n):\n",
    "                    if graph[tp][i]==1 and not find[i]:\n",
    "                        q.append(i)\n",
    "                        find[i]=True\n",
    "            #print()\n",
    "            return c\n",
    "        initial.sort()\n",
    "        ans=float(\"inf\")\n",
    "        ret=-1\n",
    "        for b in initial:\n",
    "            now=dfs(b)\n",
    "            if now<ans:\n",
    "                ans=now\n",
    "                ret=b\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        def find(x):\n",
    "            if p[x] != x:\n",
    "                p[x] = find(p[x])\n",
    "            return p[x]\n",
    "\n",
    "        def union(x, y):\n",
    "            p[find(x)] = find(y)\n",
    "\n",
    "        n = len(graph)\n",
    "        p, initial = list(range(n)), set(initial)\n",
    "        \n",
    "        # 未感染且相邻的连接在一起\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if i != j and i not in initial and j not in initial and graph[i][j] == 1:\n",
    "                    union(i, j)\n",
    "\n",
    "        # 计算每个节点集合的数量\n",
    "        sz = [0] * n\n",
    "        for i in range(n):\n",
    "            sz[find(i)] += 1\n",
    "\n",
    "        # 存储每个节点集合能被多少个感染节点所连接\n",
    "        nxt = collections.defaultdict(set)\n",
    "        for i in initial:\n",
    "            for j in range(n):\n",
    "                if graph[i][j] and j not in initial:\n",
    "                    nxt[find(j)].add(i)\n",
    "                    nxt[i].add(find(j))\n",
    "        # 如果节点集合值为1,说明去除本节点后这个集合不会被感染,进行累加计数， 因此按索引遍历，找到使得减少数量最大的感染节点即可。\n",
    "        return max(sorted(initial), key=lambda i: sum(sz[j] for j in nxt[i] if len(nxt[j]) == 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n=len(graph)\n",
    "        def dfs(x):\n",
    "            find=[False]*n\n",
    "            q=[]\n",
    "            for i in initial:\n",
    "                find[i]=True\n",
    "                if i==x:\n",
    "                    continue\n",
    "                else:\n",
    "                    q.append(i)\n",
    "                    \n",
    "            c=0\n",
    "            while q:\n",
    "                tp=q.pop()\n",
    "                #if find[tp]:\n",
    "                #    continue\n",
    "                c+=1\n",
    "                #find[tp]=True\n",
    "                for i in range(n):\n",
    "                    if graph[tp][i]==1 and not find[i]:\n",
    "                        q.append(i)\n",
    "                        find[i]=True\n",
    "            print()\n",
    "            return c\n",
    "        initial.sort()\n",
    "        ans=float(\"inf\")\n",
    "        ret=-1\n",
    "        for b in initial:\n",
    "            now=dfs(b)\n",
    "            if now<ans:\n",
    "                ans=now\n",
    "                ret=b\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        fa = list(range(n))\n",
    "        sz, szbad, bad = [1] * n, [0] * n, [0] * n\n",
    "        for i in initial: szbad[i], bad[i] = 1, 1\n",
    "        def getfa(x):\n",
    "            root = x\n",
    "            while root != fa[root]:\n",
    "                root = fa[root]\n",
    "            while x != fa[x]:\n",
    "                fa[x], x = root, fa[x] \n",
    "            return x\n",
    "\n",
    "        def merge(a, b):\n",
    "            pa, pb = getfa(a), getfa(b)\n",
    "            if pa != pb:\n",
    "                if sz[pa] > sz[pb]: pa, pb = pb, pa\n",
    "                sz[pb] += sz[pa]\n",
    "                sz[pa] = 0\n",
    "                szbad[pb] += szbad[pa]\n",
    "                szbad[pa] = 0\n",
    "                fa[pa] = pb\n",
    "        d = {}\n",
    "        g = [[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            for j in range(0, i + 1):\n",
    "                if graph[i][j] == 1:\n",
    "                    g[i].append(j)\n",
    "                    g[j].append(i)\n",
    "                    if not bad[i] and not bad[j]:\n",
    "                        merge(i, j)\n",
    "        groupGood = set()\n",
    "        for i in range(n):\n",
    "            groupGood.add(getfa(i))\n",
    "        for i in initial:\n",
    "            s = set()\n",
    "            for j in g[i]:\n",
    "                pj = getfa(j)\n",
    "                if pj not in s:\n",
    "                    szbad[pj] += 1\n",
    "                    s.add(pj)\n",
    "        \n",
    "        t = [0] * n\n",
    "        evil = set(initial)\n",
    "        for i in initial:\n",
    "            s = set()\n",
    "            for j in g[i]:\n",
    "                if j in evil: continue\n",
    "                pj = getfa(j)\n",
    "                if pj not in s and szbad[pj] == 1:\n",
    "                    t[i] += sz[pj]\n",
    "                    s.add(pj)\n",
    "        ans = initial[-1]\n",
    "        M = 0\n",
    "        for i in sorted(initial, reverse=True):\n",
    "            if t[i] > M:\n",
    "                M, ans = t[i], i\n",
    "            elif t[i] == M:\n",
    "                ans = i\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        ans = []\n",
    "        e = defaultdict(set)\n",
    "        for i in range(n):\n",
    "            for j in range(i + 1,n):\n",
    "                if graph[i][j]:\n",
    "                    e[i].add(j)\n",
    "                    e[j].add(i)\n",
    "                    \n",
    "        for i in initial:\n",
    "            vis = set()\n",
    "            q = deque()\n",
    "            for j in initial:\n",
    "                if j == i:\n",
    "                    continue\n",
    "                q.append(j)\n",
    "                vis.add(j)\n",
    "            while q:\n",
    "                cur = q.popleft()\n",
    "                for nxt in e[cur]:\n",
    "                    if nxt == i or nxt in vis:\n",
    "                        continue\n",
    "                    vis.add(nxt)\n",
    "                    q.append(nxt)\n",
    "            ans.append([len(vis),i])\n",
    "        ans.sort()\n",
    "        return ans[0][1]\n",
    "                    \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        def find(x):\n",
    "            if p[x] != x:\n",
    "                p[x] = find(p[x])\n",
    "            return p[x]\n",
    "\n",
    "        def union(x, y):\n",
    "            p[find(x)] = find(y)\n",
    "\n",
    "        n = len(graph)\n",
    "        p, initial = list(range(n)), set(initial)\n",
    "        \n",
    "        # 未感染且相邻的连接在一起\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if i != j and i not in initial and j not in initial and graph[i][j] == 1:\n",
    "                    union(i, j)\n",
    "\n",
    "        # 计算每个节点集合的数量\n",
    "        sz = [0] * n\n",
    "        for i in range(n):\n",
    "            sz[find(i)] += 1\n",
    "\n",
    "        # 存储每个节点集合能被多少个感染节点所连接\n",
    "        nxt = collections.defaultdict(set)\n",
    "        for i in initial:\n",
    "            for j in range(n):\n",
    "                if graph[i][j] and j not in initial:\n",
    "                    nxt[find(j)].add(i)\n",
    "                    nxt[i].add(find(j))\n",
    "        # 如果节点集合值为1,说明去除本节点后这个集合不会被感染,进行累加计数， 因此按索引遍历，找到使得减少数量最大的感染节点即可。\n",
    "        return max(sorted(initial), key=lambda i: sum(sz[j] for j in nxt[i] if len(nxt[j]) == 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "\n",
    "        def find(x:int) -> int:\n",
    "            if f[x] != x:\n",
    "                f[x] = find(f[x])\n",
    "            return f[x]\n",
    "        \n",
    "        def union(x, y):\n",
    "            f[find(x)] = find(y)\n",
    "\n",
    "        n = len(graph)\n",
    "        f, initial = list(range(n)), set(initial)\n",
    "        for i, j in product(range(n), range(n)):\n",
    "            if graph[i][j] and not initial & {i, j}:\n",
    "                union(i, j)\n",
    "        sz = [0] * n \n",
    "        for i in range(n):\n",
    "            sz[find(i)] += 1\n",
    "        nxt = defaultdict(set)\n",
    "        for i in initial:\n",
    "            for j in range(n):\n",
    "                if graph[i][j] and j not in initial:\n",
    "                    nxt[find(j)].add(i)\n",
    "                    nxt[i].add(find(j))\n",
    "        return max(sorted(initial), key = lambda i: sum(sz[j] for j in nxt[i] if len(nxt[j]) == 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parents = [i for i in range(n)]\n",
    "        self.sizes = [1]*n\n",
    "    def find(self, p):\n",
    "        if self.parents[p] != p:\n",
    "            self.parents[p] = self.find(self.parents[p])\n",
    "            \n",
    "        return self.parents[p]\n",
    "    def connected(self, p, q):\n",
    "        return self.find(p) == self.find(q)\n",
    "    def union(self, p, q):\n",
    "        if not self.connected(p, q):\n",
    "            p_p = self.find(p)\n",
    "            p_q = self.find(q)\n",
    "            self.parents[p_p] = p_q\n",
    "            self.sizes[p_q] += self.sizes[p_p]\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial.sort()\n",
    "        size = -1\n",
    "        ans = None\n",
    "        for node in initial:\n",
    "            uf = UnionFind(len(graph))\n",
    "            for i in range(len(graph)-1):\n",
    "                if i == node:\n",
    "                    continue\n",
    "                for j in range(i+1, len(graph)):\n",
    "                    if j == node:\n",
    "                        continue\n",
    "                    if graph[i][j]:\n",
    "                        uf.union(i, j)\n",
    "            seen = set()\n",
    "            total = 0\n",
    "            for i in initial:\n",
    "                if i == node:\n",
    "                    continue\n",
    "                root = uf.find(i)\n",
    "                if root in seen:\n",
    "                    continue\n",
    "                seen.add(root)\n",
    "                total += uf.sizes[root]\n",
    "            if len(graph)-total > size:\n",
    "                size = len(graph)-total\n",
    "                ans = node\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        self.virus = [False]*301\n",
    "        self.cnts = [0]*301\n",
    "        self.infect = [0]*301\n",
    "        self.father = [0]*301\n",
    "        self.size = [0]*301\n",
    "        def build(n,initial):\n",
    "            for i in range(n):\n",
    "                self.virus[i] = False\n",
    "                self.cnts[i] = 0\n",
    "                self.infect[i] = -1\n",
    "                self.size[i] = 1\n",
    "                self.father[i] = i\n",
    "            for i in initial:\n",
    "                self.virus[i] = True\n",
    "        def find(i):\n",
    "            if i!=self.father[i]:\n",
    "                self.father[i] = find(self.father[i])\n",
    "            return self.father[i]\n",
    "        def union(x,y):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if fx!=fy:\n",
    "                self.father[fx] = fy\n",
    "                self.size[fy] += self.size[fx]\n",
    "        n = len(graph)\n",
    "        build(n,initial)\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j]==1 and not self.virus[i] and not self.virus[j]:\n",
    "                    union(i,j)\n",
    "        for sick in initial:\n",
    "            for neighbor in range(n):\n",
    "                if sick!=neighbor and not self.virus[neighbor] and graph[sick][neighbor]==1:\n",
    "                    fn = find(neighbor)\n",
    "                    if self.infect[fn] == -1:\n",
    "                        self.infect[fn] = sick\n",
    "                    elif self.infect[fn]!=-2 and self.infect[fn]!=sick:\n",
    "                        self.infect[fn] = -2\n",
    "        for i in range(n):\n",
    "            if i==find(i) and self.infect[i]>=0:\n",
    "                self.cnts[self.infect[i]] += self.size[i]\n",
    "        initial.sort()\n",
    "        ans = initial[0]\n",
    "        max = self.cnts[ans]\n",
    "        for i in initial:\n",
    "            if self.cnts[i]>max:\n",
    "                max = self.cnts[i]\n",
    "                ans = i\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class DSU:\n",
    "    def __init__(self, N):\n",
    "        self.p = [i for i in range(N)]\n",
    "        self.sz = [1] * N\n",
    "\n",
    "    def find(self, x):\n",
    "        while self.p[x] != x:\n",
    "            x = self.p[x]\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        xr = self.find(x)\n",
    "        yr = self.find(y)\n",
    "        if xr != yr:\n",
    "            self.p[xr] = yr\n",
    "            self.sz[yr] += self.sz[xr]\n",
    "\n",
    "    def size(self, x):\n",
    "        return self.sz[self.find(x)]\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        dsu = DSU(n)\n",
    "        initilSet = set(initial)\n",
    "\n",
    "        for i in range(n):\n",
    "            if i in initilSet:\n",
    "                continue\n",
    "            for j in range(n):\n",
    "                if j in initilSet:\n",
    "                    continue\n",
    "                if i != j and graph[i][j]:\n",
    "                    dsu.union(i, j)\n",
    "\n",
    "        # dsu now represents the components of the graph without\n",
    "        # any nodes from initial.  Let's call this graph G.\n",
    "        clean = [x for x in range(n) if x not in initial]\n",
    "        count = collections.Counter()\n",
    "        node_to_compo = {}\n",
    "        for u in initial:\n",
    "            components = set()\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    components.add(dsu.find(v))\n",
    "            node_to_compo[u] = components\n",
    "\n",
    "            # 最终查找只会被一个u 感染的component\n",
    "            for c in components:\n",
    "                count[c] += 1\n",
    "\n",
    "\n",
    "        ans = (0, min(initial))\n",
    "        for u, components in node_to_compo.items():\n",
    "            addingConnect = 0\n",
    "        \n",
    "            for c in components:\n",
    "                if count[c] == 1: # 只会被一个感染， 如果是两个打破一个， 并不会减少感染\n",
    "                    addingConnect += dsu.size(c)\n",
    "\n",
    "            if addingConnect > ans[0]:\n",
    "                ans = addingConnect, u\n",
    "            elif addingConnect == ans[0] and u < ans[1]:\n",
    "                ans = addingConnect, u\n",
    "\n",
    "        return ans[1]\n",
    "            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "\n",
    "        n = len(graph)\n",
    "        parent = list(range(n))\n",
    "        size = [1] * n\n",
    "\n",
    "        def find(x: int) -> int:\n",
    "            if parent[x] != x:\n",
    "                parent[x] = find(parent[x])\n",
    "            return parent[x]\n",
    "\n",
    "        def union(x: int, y: int) -> None:\n",
    "            x_p = find(x)\n",
    "            y_p = find(y)\n",
    "            if (x_p == y_p):\n",
    "                return None\n",
    "            parent[x_p] = y_p\n",
    "            size[y_p] += size[x_p]\n",
    "\n",
    "        initial = set(initial)\n",
    "        # 提前建图\n",
    "        g = defaultdict(set)\n",
    "        for i, row in enumerate(graph):\n",
    "            for j, x in enumerate(row):\n",
    "                if x and i != j:\n",
    "                    g[i].add(j)\n",
    "                    g[j].add(i)\n",
    "\n",
    "        for o, ps in g.items():\n",
    "            if o in initial:\n",
    "                continue\n",
    "            for p in ps:\n",
    "                if p in initial:\n",
    "                    continue\n",
    "                union(o, p)\n",
    "\n",
    "        init_to_union = defaultdict(set)\n",
    "        cnt = Counter()\n",
    "        for u in initial:\n",
    "            for v in g[u]:\n",
    "                if v in initial:\n",
    "                    continue\n",
    "                init_to_union[u].add(find(v))\n",
    "            for v in init_to_union[u]:\n",
    "                cnt[find(v)] += 1\n",
    "\n",
    "        ans = 0\n",
    "        max_v = -inf\n",
    "        for u in sorted(initial):\n",
    "            res = 0\n",
    "            for v in init_to_union[u]:\n",
    "                if cnt[v] == 1:\n",
    "                    res += size[v]\n",
    "            if res > max_v:\n",
    "                ans = u\n",
    "                max_v = res\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        edges = defaultdict(list)\n",
    "        n = len(graph)\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j]:\n",
    "                    edges[i].append(j)\n",
    "                    edges[j].append(i)\n",
    "        cnt, ansNode = float(\"inf\"), None\n",
    "        for i in range(len(initial)):\n",
    "            q = initial[:i] + initial[i+1:]\n",
    "            q = deque(q)\n",
    "            curNode = initial[i]\n",
    "            visited = {node:True for node in q}\n",
    "            while q:\n",
    "                node = q.popleft()\n",
    "                for nextNode in edges[node]:\n",
    "                    if nextNode == curNode or nextNode in visited:\n",
    "                        continue\n",
    "                    q.append(nextNode)\n",
    "                    visited[nextNode] = True\n",
    "            curCnt = len(visited.keys())\n",
    "            # print(f\"curCnt:{curCnt}, curAns:{curNode}\")\n",
    "            if cnt > curCnt or (cnt == curCnt and ansNode > curNode):\n",
    "                cnt = curCnt\n",
    "                ansNode = curNode\n",
    "            # print(f\"cnt:{cnt}, ansNode:{ansNode}\")\n",
    "        return ansNode\n",
    "\n",
    "                    \n",
    "                    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        uf = UnionFind(n)\n",
    "        initial_set = set(initial)\n",
    "        clean = [x for x in range(n) if x not in initial_set]\n",
    "\n",
    "        for u in clean:\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    uf.union(u,v)\n",
    "\n",
    "        # 记录初始感染节点的联通分量\n",
    "        infect_nodes = {}\n",
    "        # 记录每个clean节点能够被多少个不同initial节点感染\n",
    "        cnt = Counter()\n",
    "        for u in initial:\n",
    "            _set = set()\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    _set.add(uf.find(v))\n",
    "            infect_nodes[u] = _set\n",
    "            for x in _set:\n",
    "                cnt[x] += 1\n",
    "        \n",
    "        ans = (-1,None)\n",
    "        for u,s in infect_nodes.items():\n",
    "            score = sum(uf.size[x] for x in s if cnt[x] == 1)\n",
    "            if score > ans[0] or (score == ans[0] and u < ans[1]):\n",
    "                ans = (score,u)\n",
    "        return ans[1]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self,n):\n",
    "        self.fa = list(range(n))\n",
    "        self.size = [1]*n\n",
    "    \n",
    "    def find(self,x):\n",
    "        if self.fa[x] != x:\n",
    "            self.fa[x] = self.find(self.fa[x])\n",
    "        return self.fa[x]\n",
    "\n",
    "    def union(self,x,y):\n",
    "        a,b = self.find(x),self.find(y)\n",
    "        if a != b:\n",
    "            self.fa[a] = b\n",
    "            self.size[b] += self.size[a]\n",
    "    \n",
    "    def get_size(self,x):\n",
    "        return self.size[self.find(x)]\n",
    "    \n",
    "    def is_conn(self,x,y):\n",
    "        return self.find(x) == self.find(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:        \n",
    "        n = len(graph)\n",
    "        edge = defaultdict(dict)\n",
    "        for i in range(n):\n",
    "            for j in range(i+1, n):\n",
    "                if graph[i][j]:\n",
    "                    edge[i][j] = 1\n",
    "                    edge[j][i] = 1 \n",
    "        # 初始化\n",
    "        m = len(initial)\n",
    "        ans = n\n",
    "        cnt = float('inf')\n",
    "        visit = dict()\n",
    "        for node in initial:\n",
    "            visit[node] = 1\n",
    "        # 枚举\n",
    "        for i in range(m):\n",
    "            stack = initial[:]\n",
    "            k = stack.pop(i)\n",
    "            # 感染记录\n",
    "            visit_cur = visit.copy()\n",
    "            del visit_cur[k]  \n",
    "            # 删除节点关系\n",
    "            nodes = list(edge[k].keys())\n",
    "            for node in edge[k]:\n",
    "                del edge[node][k]\n",
    "            del edge[k] \n",
    "            # 广搜感染\n",
    "            while stack:\n",
    "                nex = []\n",
    "                for cur in stack:\n",
    "                    for node in edge[cur]:\n",
    "                        if node not in visit_cur:\n",
    "                            nex.append(node)\n",
    "                            visit_cur[node] = 1\n",
    "                stack = nex[:] \n",
    "            # 比较\n",
    "            if len(visit_cur) < cnt or (len(visit_cur) == cnt and k < ans):\n",
    "                ans = k\n",
    "                cnt = len(visit_cur)  \n",
    "            # 还原图关系\n",
    "            for node in nodes:\n",
    "                edge[node][k] = 1\n",
    "                edge[k][node] = 1\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        uf = UnionFind(n)\n",
    "        initial_set = set(initial)\n",
    "        clean = [x for x in range(n) if x not in initial_set]\n",
    "\n",
    "        for u in clean:\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    uf.union(u,v)\n",
    "\n",
    "        # 记录初始感染节点的联通分量\n",
    "        infect_nodes = {}\n",
    "        # 记录每个clean节点能够被多少个不同initial节点感染\n",
    "        cnt = Counter()\n",
    "        for u in initial:\n",
    "            _set = set()\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    _set.add(uf.find(v))\n",
    "            infect_nodes[u] = _set\n",
    "            for x in _set:\n",
    "                cnt[x] += 1\n",
    "        \n",
    "        ans = (-1,None)\n",
    "        for u,s in infect_nodes.items():\n",
    "            score = sum(uf.size[x] for x in s if cnt[x] == 1)\n",
    "            if score > ans[0] or (score == ans[0] and u < ans[1]):\n",
    "                ans = (score,u)\n",
    "        return ans[1]\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self,n):\n",
    "        self.fa = list(range(n))\n",
    "        self.size = [1]*n\n",
    "    \n",
    "    def find(self,x):\n",
    "        if self.fa[x] != x:\n",
    "            self.fa[x] = self.find(self.fa[x])\n",
    "        return self.fa[x]\n",
    "\n",
    "    def union(self,x,y):\n",
    "        a,b = self.find(x),self.find(y)\n",
    "        if a != b:\n",
    "            self.fa[a] = b\n",
    "            self.size[b] += self.size[a]\n",
    "    \n",
    "    def get_size(self,x):\n",
    "        return self.size[self.find(x)]\n",
    "    \n",
    "    def is_conn(self,x,y):\n",
    "        return self.find(x) == self.find(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initials=set(initial)\n",
    "        initial_dict=collections.defaultdict(set)\n",
    "        n=len(graph)\n",
    "        fa=[i for i in range(n)]\n",
    "        def find(i):\n",
    "            if fa[i]!=i:\n",
    "                fa[i]=find(fa[i])\n",
    "            return fa[i]\n",
    "        for i,g in enumerate(graph):\n",
    "            for j,ge in enumerate(g):\n",
    "                if i in initials: \n",
    "                    if ge and j!=i and j not in initials:\n",
    "                        initial_dict[i].add(j)\n",
    "                elif ge:\n",
    "                    #print(i,ge,initials,ge not in initials)\n",
    "                    if j not in initials:\n",
    "                        fi,fj=find(i),find(j)\n",
    "                        #print('notin',fi,j)\n",
    "                        if fi!=fj:\n",
    "                            print(fi,fj)\n",
    "                            fa[fi]=fj\n",
    "        print(fa)\n",
    "        sizes=[0]*n\n",
    "        for i in range(n):\n",
    "            sizes[find(i)]+=1\n",
    "        \n",
    "        dirty_source=[-1]*n   \n",
    "        for k,v in initial_dict.items():\n",
    "             for ve in v:\n",
    "                 if dirty_source[find(ve)]==-1 or dirty_source[find(ve)]==k:\n",
    "                     dirty_source[find(ve)]=k     \n",
    "                 else:\n",
    "                     dirty_source[find(ve)]=-2\n",
    "        print(sizes,dirty_source,initial_dict)\n",
    "        ms=0;ans=-1\n",
    "        cnt=[0]*n\n",
    "        for s,d in zip(sizes,dirty_source):\n",
    "            if d>-1:\n",
    "                cnt[d]+=s    \n",
    "        return min(initial) if all([e==0 for e in cnt]) else cnt.index(max(cnt))\n",
    "\n",
    "\n",
    "\n",
    "         \n",
    "\n",
    "\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        uf = UnionFind(n)\n",
    "        initial_set = set(initial)\n",
    "        clean = [x for x in range(n) if x not in initial_set]\n",
    "\n",
    "        for u in clean:\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    uf.union(u,v)\n",
    "\n",
    "        # 记录初始感染节点的联通分量\n",
    "        infect_nodes = {}\n",
    "        # 记录每个clean节点能够被多少个不同initial节点感染\n",
    "        cnt = Counter()\n",
    "        for u in initial:\n",
    "            _set = set()\n",
    "            for v in clean:\n",
    "                if graph[u][v]:\n",
    "                    _set.add(uf.find(v))\n",
    "            infect_nodes[u] = _set\n",
    "            for x in _set:\n",
    "                cnt[x] += 1\n",
    "        \n",
    "        ans = (-1,None)\n",
    "        for u,s in infect_nodes.items():\n",
    "            print(u,s)\n",
    "            score = sum(uf.size[x] for x in s if cnt[x] == 1)\n",
    "            if score > ans[0] or (score == ans[0] and u < ans[1]):\n",
    "                ans = (score,u)\n",
    "        return ans[1]\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self,n):\n",
    "        self.fa = list(range(n))\n",
    "        self.size = [1]*n\n",
    "    \n",
    "    def find(self,x):\n",
    "        if self.fa[x] != x:\n",
    "            self.fa[x] = self.find(self.fa[x])\n",
    "        return self.fa[x]\n",
    "\n",
    "    def union(self,x,y):\n",
    "        a,b = self.find(x),self.find(y)\n",
    "        if a != b:\n",
    "            self.fa[a] = b\n",
    "            self.size[b] += self.size[a]\n",
    "    \n",
    "    def get_size(self,x):\n",
    "        return self.size[self.find(x)]\n",
    "    \n",
    "    def is_conn(self,x,y):\n",
    "        return self.find(x) == self.find(y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        d = defaultdict(list)\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j] == 1:\n",
    "                    d[i].append(j)\n",
    "        initial = sorted(initial)\n",
    "        l = len(initial)\n",
    "        M = [n, initial[0]]\n",
    "        for i in range(l):\n",
    "            q = deque(initial[:i] + initial[i + 1:])\n",
    "            p = []\n",
    "            while q:\n",
    "                temp = q.popleft()\n",
    "                if temp not in p:\n",
    "                    p.append(temp)\n",
    "                    for x in d[temp]:\n",
    "                        if x != initial[i] and x not in p:\n",
    "                            q.append(x)\n",
    "            if len(p) < M[0]:\n",
    "                M = [len(p), initial[i]]\n",
    "        return M[1]\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\r\n",
    "        initial = set(initial)\r\n",
    "        def find(x):#并查集\r\n",
    "            if x not in cache:\r\n",
    "                cache[x] = x\r\n",
    "            if cache[x] != x:\r\n",
    "                cache[x] = find(cache[x])\r\n",
    "            return cache[x]\r\n",
    "        def merge(x, y):\r\n",
    "            cache[find(x)] = find(y)\r\n",
    "        n = len(graph)\r\n",
    "        cache = {}\r\n",
    "        for i in range(n):\r\n",
    "            for j in range(n):\r\n",
    "                if i != j and i not in initial and j not in initial and graph[i][j] == 1:\r\n",
    "                    merge(i, j)\r\n",
    "        d = {}#计算每个节点集合的数量\r\n",
    "        for i in range(n):\r\n",
    "            i = find(i)\r\n",
    "            d[i] = d.get(i, 0) + 1\r\n",
    "        s = set(d.keys())\r\n",
    "        cnt = {}#存储每个节点集合能被多少个感染节点所连接\r\n",
    "        for i in initial:\r\n",
    "            c = 0\r\n",
    "            s1 = set([find(j) for j in range(n) if j not in initial and graph[i][j] == 1])#计算的是感染节点数,避免重复计数采取集合\r\n",
    "            for j in s1:\r\n",
    "                cnt[j] = cnt.get(j, 0) + 1\r\n",
    "        m = -1\r\n",
    "        for i in sorted(initial):#升序遍历,如果有多个节点需要返回最小的\r\n",
    "            c = 0\r\n",
    "            s1 = set([find(j) for j in range(n) if j not in initial and graph[i][j] == 1])\r\n",
    "            for j in s1:\r\n",
    "                if cnt[j] == 1:#如果节点集合值为1,说明去除本节点后这个集合不会被感染,进行累加计数\r\n",
    "                    c += d[j]\r\n",
    "            if c > m:#当不会被感染的集合数量越大说明被感染节点总数越小\r\n",
    "                m = c\r\n",
    "                res = i\r\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import copy\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        min_sum = n+1\n",
    "        min_node = -1\n",
    "        visited = set()\n",
    "        g = defaultdict(list)\n",
    "\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j] == 1:\n",
    "                    g[i].append(j)\n",
    "\n",
    "        initial.sort()\n",
    "        for i in range(len(initial)):\n",
    "            queue = initial[:i] + initial[i+1:]\n",
    "            del_node = initial[i]\n",
    "            visited.clear()\n",
    "\n",
    "            # cg = copy.deepcopy(g)\n",
    "            # cg.pop(del_node)\n",
    "            # for k,v in cg.items():\n",
    "            #     for \n",
    "\n",
    "            for j in range(len(queue)):\n",
    "                visited.add(queue[j])\n",
    "\n",
    "            while queue:\n",
    "                pop_node = queue[0]\n",
    "                queue.pop(0)\n",
    "                for nextN in g[pop_node]:\n",
    "                    if nextN not in visited and nextN!=del_node:\n",
    "                        queue.append(nextN)\n",
    "                        visited.add(nextN)\n",
    "            \n",
    "            cur_sum = len(visited)\n",
    "            if cur_sum<min_sum:\n",
    "                min_sum = cur_sum\n",
    "                min_node = initial[i]\n",
    "        \n",
    "        return min_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from collections import defaultdict\n",
    "import copy\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        min_sum = n+1\n",
    "        min_node = -1\n",
    "        visited = set()\n",
    "        g = defaultdict(list)\n",
    "\n",
    "        for i in range(n):\n",
    "            for j in range(n):\n",
    "                if graph[i][j] == 1:\n",
    "                    g[i].append(j)\n",
    "\n",
    "        initial.sort()\n",
    "        for i in range(len(initial)):\n",
    "            queue = initial[:i] + initial[i+1:]\n",
    "            del_node = initial[i]\n",
    "            visited.clear()\n",
    "\n",
    "            for j in range(len(queue)):\n",
    "                visited.add(queue[j])\n",
    "\n",
    "            while queue:\n",
    "                pop_node = queue[0]\n",
    "                queue.pop(0)\n",
    "                for nextN in g[pop_node]:\n",
    "                    if nextN not in visited and nextN!=del_node:\n",
    "                        queue.append(nextN)\n",
    "                        visited.add(nextN)\n",
    "            \n",
    "            cur_sum = len(visited)\n",
    "            if cur_sum<min_sum:\n",
    "                min_sum = cur_sum\n",
    "                min_node = initial[i]\n",
    "        \n",
    "        return min_node"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n=len(graph)\n",
    "        ad=collections.defaultdict(list)\n",
    "        for i in range(n-1):\n",
    "            for j in range(i+1,n):\n",
    "                if graph[i][j]==1:\n",
    "                    ad[i].append(j)\n",
    "                    ad[j].append(i)\n",
    "        initial.sort()\n",
    "        mi=float(\"inf\")\n",
    "        ans=-1\n",
    "        for x in initial:\n",
    "            vis=[0]*n\n",
    "            stack=initial.copy()\n",
    "            while stack:\n",
    "                cur=stack.pop()\n",
    "                if cur==x:\n",
    "                    continue\n",
    "                if vis[cur]==1:\n",
    "                    continue\n",
    "                vis[cur]=1\n",
    "                for nex in ad[cur]:\n",
    "                    if vis[x]==0 and nex!=x:\n",
    "                        stack.append(nex)\n",
    "            tmp=sum(vis)\n",
    "          #  print(vis)\n",
    "          #  print(tmp)\n",
    "            if tmp<mi:\n",
    "                ans=x\n",
    "                mi=tmp\n",
    "        return ans\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\r\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\r\n",
    "        n = len(graph)\r\n",
    "        initial.sort()\r\n",
    "        mn = n + 1 \r\n",
    "        res = -1 \r\n",
    "        g = [[] for _ in range(n)]\r\n",
    "        for i in range(n):\r\n",
    "            for j in range(n):\r\n",
    "                if graph[i][j]:\r\n",
    "                    g[i].append(j)\r\n",
    "                \r\n",
    "        for i, x in enumerate(initial):\r\n",
    "            vis = [False] * n \r\n",
    "\r\n",
    "            def dfs(j):\r\n",
    "                if j == x or vis[j]:\r\n",
    "                    return\r\n",
    "                vis[j] = True  \r\n",
    "                for v in g[j]:\r\n",
    "                    dfs(v)\r\n",
    "\r\n",
    "            for j in initial:\r\n",
    "                dfs(j)\r\n",
    "            s = sum(vis)\r\n",
    "            if s < mn:\r\n",
    "                res = i \r\n",
    "                mn = s\r\n",
    "        return initial[res]\r\n",
    "\r\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parents = [i for i in range(n)]\n",
    "        self.sizes = [1]*n\n",
    "    def find(self, p):\n",
    "        if self.parents[p] != p:\n",
    "            self.parents[p] = self.find(self.parents[p])\n",
    "            \n",
    "        return self.parents[p]\n",
    "    def connected(self, p, q):\n",
    "        return self.find(p) == self.find(q)\n",
    "    def union(self, p, q):\n",
    "        if not self.connected(p, q):\n",
    "            p_p = self.find(p)\n",
    "            p_q = self.find(q)\n",
    "            self.parents[p_p] = p_q\n",
    "            self.sizes[p_q] += self.sizes[p_p]\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial.sort()\n",
    "        size = -1\n",
    "        ans = None\n",
    "        edges = []\n",
    "        for i in range(len(graph)-1):\n",
    "            for j in range(i+1, len(graph)):\n",
    "                if graph[i][j]:\n",
    "                    edges.append([i,j])\n",
    "        for node in initial:\n",
    "            uf = UnionFind(len(graph))\n",
    "            for edge in edges:\n",
    "                if node in edge:\n",
    "                    continue\n",
    "                uf.union(*edge)\n",
    "            seen = set()\n",
    "            total = 0\n",
    "            for i in initial:\n",
    "                if i == node:\n",
    "                    continue\n",
    "                root = uf.find(i)\n",
    "                if root in seen:\n",
    "                    continue\n",
    "                seen.add(root)\n",
    "                total += uf.sizes[root]\n",
    "            if len(graph)-total > size:\n",
    "                size = len(graph)-total\n",
    "                ans = node\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        fa = list(range(n))\n",
    "        sz, szbad, bad = [1] * n, [0] * n, [0] * n\n",
    "        for i in initial: szbad[i], bad[i] = 1, 1\n",
    "        def getfa(x):\n",
    "            root = x\n",
    "            while root != fa[root]:\n",
    "                root = fa[root]\n",
    "            while x != fa[x]:\n",
    "                fa[x], x = root, fa[x] \n",
    "            return x\n",
    "\n",
    "        def merge(a, b):\n",
    "            pa, pb = getfa(a), getfa(b)\n",
    "            if pa != pb:\n",
    "                if sz[pa] > sz[pb]: pa, pb = pb, pa\n",
    "                sz[pb] += sz[pa]\n",
    "                sz[pa] = 0\n",
    "                szbad[pb] += szbad[pa]\n",
    "                szbad[pa] = 0\n",
    "                fa[pa] = pb\n",
    "        d = {}\n",
    "        g = [[] for _ in range(n)]\n",
    "        for i in range(n):\n",
    "            for j in range(0, i+1):\n",
    "                if graph[i][j] == 1:\n",
    "                    merge(i, j)\n",
    "                    g[i].append(j)\n",
    "                    g[j].append(i)\n",
    "                    # if i != j: print(i, j)\n",
    "        \n",
    "        # 割点\n",
    "        dfn, low, cut = [-1]*n, [-1]*n, [False]*n\n",
    "        idx, count = 0,0\n",
    "\n",
    "        def tarjan(x, fa):\n",
    "            nonlocal idx, count\n",
    "            child = 0\n",
    "            dfn[x], low[x] = idx, idx\n",
    "            idx+=1\n",
    "            for y in g[x]:\n",
    "                if dfn[y]==-1:\n",
    "                    child += 1\n",
    "                    tarjan(y,x)\n",
    "                    low[x] = min(low[x], low[y])\n",
    "                    if fa!=-1 and low[y]>=dfn[x] and cut[x]==False:\n",
    "                        cut[x] = True\n",
    "                        count+=1\n",
    "                else:\n",
    "                    if y!=fa:\n",
    "                        low[x] = min(low[x], dfn[y])\n",
    "            if fa==-1 and child>=2 and cut[x]==False:\n",
    "                cut[x] = True\n",
    "                count += 1\n",
    "            pass\n",
    "        \n",
    "        for i in range(n):\n",
    "            if dfn[i]==-1: tarjan(i,-1)\n",
    "\n",
    "        ans = initial[-1]\n",
    "        M = 0\n",
    "\n",
    "        # 求割点的未感染子树的节点个数和\n",
    "        def dfs(root, u, p):\n",
    "            vis[u] = 1\n",
    "            res = int(root != u)\n",
    "            if u != root and bad[u]:\n",
    "                return 0\n",
    "            f = False\n",
    "            for x in g[u]:\n",
    "                if vis[x]: continue\n",
    "                t = dfs(root, x, u)\n",
    "                if t == 0 and u != root:\n",
    "                    f = True\n",
    "                res += t\n",
    "            if f: return 0\n",
    "            return res\n",
    "\n",
    "        for i in range(n-1, -1, -1):\n",
    "            if bad[i] == 0:continue\n",
    "            t = getfa(i)\n",
    "            if szbad[t] == 1:\n",
    "                if sz[t] - 1 >= M: ans, M = i, sz[t] - 1\n",
    "            elif szbad[t] > 1:\n",
    "                if M == 0: ans = i\n",
    "                if cut[i]:\n",
    "                    vis = [0] * n\n",
    "                    num = dfs(i, i, -1)\n",
    "                    # print(i, num)\n",
    "                    if num >= M: ans, M = i, num\n",
    "            # print(M, i, t, szbad[t], sz[t], cut[i])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        arr = [0] * n \n",
    "        for x in initial:\n",
    "            arr[x] = 1 \n",
    "\n",
    "        def dfs(x, ban, vis):\n",
    "            #nonlocal tot \n",
    "            #tot += 1\n",
    "            vis[x] = 1\n",
    "            for i in range(n):\n",
    "                if i != x and i != ban and graph[x][i] == 1 and not vis[i]:\n",
    "                    dfs(i, ban, vis) \n",
    "\n",
    "        ans = float('inf')\n",
    "        idx = -1 \n",
    "        for i in sorted(initial):\n",
    "            vis = [0] * n \n",
    "            for j in range(n):\n",
    "                if i != j and arr[j] == 1 and not vis[j]:\n",
    "                    dfs(j, i, vis)\n",
    "            tot = sum(vis) \n",
    "            if tot < ans:\n",
    "                ans = tot \n",
    "                idx = i \n",
    "            \n",
    "        return idx  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        clean = set(range(n)) - set(initial)\n",
    "        def dfs(u):\n",
    "            for i in range(n):\n",
    "                if graph[u][i] == 1 and i not in seen and i in clean:\n",
    "                    seen.add(i)\n",
    "                    dfs(i)\n",
    "        infect = defaultdict(list)\n",
    "        for u in initial:\n",
    "            seen = set()\n",
    "            dfs(u)\n",
    "            for v in seen:\n",
    "                infect[v].append(u)\n",
    "        contribution = Counter()\n",
    "        for key, val in infect.items():\n",
    "            if len(val) == 1:\n",
    "                contribution[val[0]] += 1\n",
    "        \n",
    "        ans = [-1, min(initial)]\n",
    "        for key, val in contribution.items():\n",
    "            if val > ans[0] or (val == ans[0] and key < ans[1]):\n",
    "                ans = [val, key]\n",
    "        return ans[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        ans = float('inf')\n",
    "        maxVal = float('inf')\n",
    "        for x in initial:\n",
    "            def dfs(u):\n",
    "                if vis[u]:\n",
    "                    return 0\n",
    "                vis[u] = True\n",
    "                s = 1\n",
    "                for i in range(n):\n",
    "                    if i != x and graph[u][i] == 1 and not vis[i]:\n",
    "                        s += dfs(i)\n",
    "                return s\n",
    "            res = 0\n",
    "            vis = [False] * n\n",
    "            for u in initial:\n",
    "                if u != x:\n",
    "                    res += dfs(u)\n",
    "            if res < maxVal:\n",
    "                maxVal = res\n",
    "                ans = x\n",
    "            elif res == maxVal:\n",
    "                ans = min(ans, x)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        initial = set(initial)\n",
    "        # 提前建图\n",
    "        g = defaultdict(set)\n",
    "        for i, row in enumerate(graph):\n",
    "            for j, x in enumerate(row):\n",
    "                if x and i != j:\n",
    "                    g[i].add(j)\n",
    "                    g[j].add(i)\n",
    "\n",
    "        def dfs(o: int):\n",
    "            '''\n",
    "            从节点o开始能感染到的节点，dfs的顶层入口是感染源\n",
    "            假设只有一个感染源，那对于其他感染源，要当他们不存在\n",
    "            \n",
    "            图中存在环\n",
    "            '''\n",
    "            for p in g[o]:\n",
    "                if p in v or p in initial:\n",
    "                    continue\n",
    "                v.add(p)\n",
    "                dfs(p)\n",
    "        \n",
    "        # 每个干净节点会被哪些感染源感染\n",
    "        infected = defaultdict(list)\n",
    "        cnt_clean = Counter()\n",
    "        for o in initial:\n",
    "            v = set()\n",
    "            dfs(o)\n",
    "            for clean in v:\n",
    "                infected[clean].append(o)\n",
    "                cnt_clean[clean] += 1\n",
    "                \n",
    "        # 如果某个干净节点只被一个感染源感染，那去掉这个感染源，就可以避免被感染\n",
    "        # 一个感染源，如果，“只能被它感染的节点” 越多，移除后的收益越大\n",
    "        init_cnt = Counter()\n",
    "        for clean, cnt in cnt_clean.items():\n",
    "            if cnt == 1:\n",
    "                init_cnt[infected[clean][0]] += 1\n",
    "        # 选择收益最大的节点\n",
    "        ans = 0\n",
    "        max_cnt = -inf\n",
    "        for o in sorted(initial):\n",
    "            if init_cnt[o] > max_cnt:\n",
    "                max_cnt = init_cnt[o]\n",
    "                ans = o\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        N = len(graph)\n",
    "        clean = set(range(N)) - set(initial)\n",
    "        def dfs(u, seen):\n",
    "            for v, adj in enumerate(graph[u]):\n",
    "                if adj and v in clean and v not in seen:\n",
    "                    seen.add(v)\n",
    "                    dfs(v, seen)\n",
    "\n",
    "        # For each node u in initial, dfs to find\n",
    "        # 'seen': all nodes not in initial that it can reach.\n",
    "        infected_by = {v: [] for v in clean}\n",
    "        for u in initial:\n",
    "            seen = set()\n",
    "            dfs(u, seen)\n",
    "\n",
    "            # For each node v that was seen, u infects v.\n",
    "            for v in seen:\n",
    "                infected_by[v].append(u)\n",
    "\n",
    "        # For each node u in initial, for every v not in initial\n",
    "        # that is uniquely infected by u, add 1 to the contribution for u.\n",
    "        contribution = collections.Counter()\n",
    "        for v, neighbors in infected_by.items():\n",
    "            if len(neighbors) == 1:\n",
    "                contribution[neighbors[0]] += 1\n",
    "\n",
    "        # Take the best answer.\n",
    "        best = (-1, min(initial))\n",
    "        for u, score in contribution.items():\n",
    "            if score > best[0] or score == best[0] and u < best[1]:\n",
    "                best = score, u\n",
    "        return best[1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "import collections\n",
    "\n",
    "class Solution(object):\n",
    "    def minMalwareSpread(self, graph, initial):\n",
    "        N = len(graph)\n",
    "        clean = set(range(N)) - set(initial)\n",
    "\n",
    "        def dfs(u, seen):\n",
    "            for v, adj in enumerate(graph[u]):\n",
    "                if adj and v in clean and v not in seen:\n",
    "                    seen.add(v)\n",
    "                    dfs(v, seen)\n",
    "\n",
    "        infected_by = {v: [] for v in clean}\n",
    "        for u in initial:\n",
    "            seen = set()\n",
    "            dfs(u, seen)\n",
    "\n",
    "            for v in seen:\n",
    "                infected_by[v].append(u)\n",
    "\n",
    "        contribution = collections.Counter()\n",
    "        for v, neighbors in infected_by.items():\n",
    "            if len(neighbors) == 1:\n",
    "                contribution[neighbors[0]] += 1\n",
    "\n",
    "        best = (-1, min(initial))\n",
    "        for u, score in contribution.items():\n",
    "            if score > best[0] or (score == best[0] and u < best[1]):\n",
    "                best = (score, u)\n",
    "\n",
    "        return best[1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        # Three types of node:\n",
    "        # -1   -> not visit\n",
    "        # i    -> first pass by node i\n",
    "        # [0 ~ n-1] -> another infected node\n",
    "        n = len(graph)\n",
    "        vis = [-1] * n\n",
    "        for i in initial:\n",
    "            vis[i] = i\n",
    "        \n",
    "        def dfs(node, cur):\n",
    "            vis[node] = cur \n",
    "            cnt = 1\n",
    "            for i, v in enumerate(graph[node]):\n",
    "                if not v:  continue\n",
    "                if vis[i] == -1:\n",
    "                    cnt += dfs(i, cur)\n",
    "                elif vis[i] != cur:\n",
    "                    # The set is infected by another node\n",
    "                    vis[cur] = n\n",
    "            return cnt\n",
    "        \n",
    "        ans = [min(initial), 0]\n",
    "        for i in sorted(initial):\n",
    "            count = 0\n",
    "            for j, v in enumerate(graph[i]):\n",
    "                if not v or vis[j] != -1: \n",
    "                    continue\n",
    "                # initialize\n",
    "                vis[i] = i\n",
    "                k = dfs(j, i)\n",
    "                if vis[i] == i:\n",
    "                    count += k\n",
    "            if count > ans[1]:\n",
    "                ans = [i, count]\n",
    "\n",
    "        return ans[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        vis = defaultdict(bool)\n",
    "        tree = defaultdict(set)\n",
    "        root = -1\n",
    "        initial = set(initial)\n",
    "        cnt = Counter()\n",
    "\n",
    "        def dfs(a):\n",
    "            tree[root].add(a)\n",
    "            if a in initial:\n",
    "                return\n",
    "            vis[a] = True\n",
    "            for b, t in enumerate(graph[a]):\n",
    "                if t and not vis[b]:\n",
    "                    dfs(b)\n",
    "\n",
    "\n",
    "        for a in range(n):\n",
    "            if not vis[a] and a not in initial:\n",
    "                root = a\n",
    "                dfs(a)\n",
    "                infect = tree[root] & initial\n",
    "                if len(infect) == 1:\n",
    "                    cnt[infect.pop()] += len(tree[root]) - 1\n",
    "\n",
    "        initial = list(initial)\n",
    "        # print(tree)\n",
    "        # print(cnt)\n",
    "        initial.sort()\n",
    "        ans = max_infect = -1\n",
    "        for a in initial:\n",
    "            if cnt[a] > max_infect:\n",
    "                max_infect = cnt[a]\n",
    "                ans = a\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        N = len(graph)\n",
    "        clean = set(range(N)) - set(initial)\n",
    "        def dfs(u: int, seen: set) -> None:\n",
    "            for v, adj in enumerate(graph[u]):\n",
    "                if adj and v in clean and v not in seen:\n",
    "                    seen.add(v)\n",
    "                    dfs(v, seen)\n",
    "\n",
    "        # For each node u in initial, dfs to find\n",
    "        # 'seen': all nodes not in initial that it can reach.\n",
    "        infected_by = {v: [] for v in clean}\n",
    "        for u in initial:\n",
    "            seen = set()\n",
    "            dfs(u, seen)\n",
    "\n",
    "            # For each node v that was seen, u infects v.\n",
    "            for v in seen:\n",
    "                infected_by[v].append(u)\n",
    "\n",
    "        # For each node u in initial, for every v not in initial\n",
    "        # that is uniquely infected by u, add 1 to the contribution for u.\n",
    "        contribution = collections.Counter()\n",
    "        for v, neighbors in infected_by.items():\n",
    "            if len(neighbors) == 1:\n",
    "                contribution[neighbors[0]] += 1\n",
    "\n",
    "        # Take the best answer.\n",
    "        best = (-1, min(initial))\n",
    "        for u, score in contribution.items():\n",
    "            if score > best[0] or score == best[0] and u < best[1]:\n",
    "                best = score, u\n",
    "        return best[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution(object):\n",
    "    def minMalwareSpread(self, graph, initial):\n",
    "        N = len(graph)\n",
    "        clean = set(range(N)) - set(initial)\n",
    "        def dfs(u, seen):\n",
    "            for v, adj in enumerate(graph[u]):\n",
    "                if adj and v in clean and v not in seen:\n",
    "                    seen.add(v)\n",
    "                    dfs(v, seen)\n",
    "\n",
    "        # For each node u in initial, dfs to find\n",
    "        # 'seen': all nodes not in initial that it can reach.\n",
    "        infected_by = {v: [] for v in clean}\n",
    "        for u in initial:\n",
    "            seen = set()\n",
    "            dfs(u, seen)\n",
    "\n",
    "            # For each node v that was seen, u infects v.\n",
    "            for v in seen:\n",
    "                infected_by[v].append(u)\n",
    "\n",
    "        # For each node u in initial, for every v not in initial\n",
    "        # that is uniquely infected by u, add 1 to the contribution for u.\n",
    "        contribution = Counter()\n",
    "        for v, neighbors in infected_by.items():\n",
    "            if len(neighbors) == 1:\n",
    "                contribution[neighbors[0]] += 1\n",
    "\n",
    "        # Take the best answer.\n",
    "        best = (-1, min(initial))\n",
    "        for u, score in contribution.items():\n",
    "            if score > best[0] or score == best[0] and u < best[1]:\n",
    "                best = score, u\n",
    "        return best[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def minMalwareSpread(self, graph: List[List[int]], initial: List[int]) -> int:\n",
    "        n = len(graph)\n",
    "        vis = defaultdict(bool)\n",
    "        infection = set()\n",
    "        initial = set(initial)\n",
    "\n",
    "\n",
    "        def dfs1(a):\n",
    "            vis[a] = True\n",
    "            if a in initial or a in infection:\n",
    "                return 1\n",
    "            res = 0\n",
    "            for b, t in enumerate(graph[a]):\n",
    "                if not vis[b] and t:\n",
    "                    res += dfs1(b)\n",
    "            return res\n",
    "                    \n",
    "\n",
    "\n",
    "        for a in range(n):\n",
    "            if a not in initial:\n",
    "                if dfs1(a) > 1:\n",
    "                    infection.add(a)\n",
    "                vis.clear()\n",
    "\n",
    "\n",
    "        def dfs2(a):\n",
    "            vis[a] = True\n",
    "            res = 0\n",
    "            for b, t in enumerate(graph[a]):\n",
    "                if t and not vis[b] and b not in infection and b not in initial:\n",
    "                    res += dfs2(b) + 1\n",
    "            return res\n",
    "\n",
    "        ans = -1\n",
    "        max_infect = 0\n",
    "\n",
    "        for a in initial:\n",
    "            cur_infect = dfs2(a)\n",
    "            if cur_infect > max_infect:\n",
    "                ans = a\n",
    "                max_infect = cur_infect\n",
    "            elif cur_infect == max_infect and a < ans:\n",
    "                ans = a\n",
    "\n",
    "\n",
    "        return ans if max_infect > 0 else sorted(list(initial))[0]"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
